#!/usr/bin/env python3

import argparse
import numpy as np
import xarray as xr
from multiprocessing import Pool
from itertools import zip_longest, repeat
import metpy.calc as calc
from metpy.units import units
import metpy.constants as const
import pandas as pd
import pendulum
from glob import glob
import os
import sys
sys.path.append(f'{os.path.dirname(os.path.realpath(__file__))}/../../utils')
from utils import cli, parse_time, parse_config, create_kdtree, find_bbox, create_bilinear_wgts, apply_bilinear

def sim_synop_era5(file_root, sid, lon, lat, start_time, end_time):
	cli.notice(f'Simulate SYNOP at {sid}@({lon},{lat}).')
	files = file_root + '/era5_*_sfc.grb'
	used_files = []
	times = []
	t0 = np.datetime64('1970-01-01T00:00:00')
	for file in sorted(glob(files)):
		f = xr.open_dataset(file, engine='cfgrib', backend_kwargs={'filter_by_keys':{'typeOfLevel':'surface'}})
		time = pendulum.from_timestamp(float((f.time - t0) / np.timedelta64(1, 's')))
		if start_time <= time <= end_time:
			times.append(time)
			used_files.append(file)
	if len(used_files) == 0:
		cli.error(f'Cannot find valid wrfout files!')
	f = xr.open_mfdataset(used_files, engine='cfgrib', backend_kwargs={'filter_by_keys':{'typeOfLevel':'surface'}}, concat_dim='time', combine='nested')

	zs   = calc.geopotential_to_height(f.z.interp(longitude=lon, latitude=lat, method='linear') * units('m**2/s**2'))
	u10  = f.u10.interp(longitude=lon, latitude=lat, method='linear') * units.mps
	v10  = f.v10.interp(longitude=lon, latitude=lat, method='linear') * units.mps
	t2   = f.t2m.interp(longitude=lon, latitude=lat, method='linear') * units.K
	d2   = f.d2m.interp(longitude=lon, latitude=lat, method='linear') * units.K
	ps   = f.sp .interp(longitude=lon, latitude=lat, method='linear') * units.Pa
	slp  = f.msl.interp(longitude=lon, latitude=lat, method='linear') * units.Pa
	ws10 = calc.wind_speed(u10, v10)
	wd10 = calc.wind_direction(u10, v10)
	rh2  = calc.relative_humidity_from_dewpoint(t2, d2) * 100 * units.percent
	sh2  = calc.specific_humidity_from_dewpoint(ps, d2)
	data = {
		'sid' : [sid for _ in range(len(times))],
		'lon' : [lon for _ in range(len(times))],
		'lat' : [lat for _ in range(len(times))],
		'zs'  : [float(x) for x in zs ],
		'time': times,
		'u10' : [float(x) for x in u10],
		'v10' : [float(x) for x in v10],
		't2'  : [float(x) for x in t2 ],
		'sh2' : [float(x) for x in sh2],
		'ps'  : [float(x) for x in ps ],
		'slp' : [float(x) for x in slp],
		'ws10': [float(x) for x in ws10],
		'wd10': [float(x) for x in wd10],
		'rh2' : [float(x) for x in rh2],
		'd2'  : [float(x) for x in d2 ]
	}
	return pd.DataFrame(data)

def sim_synop_ecwf(file_root, sid, lon, lat, start_time, end_time):
	cli.notice(f'Simulate SYNOP at {sid}@({lon},{lat}).')
	files = file_root + '/W_NAFP_C_ECMF_*'
	used_files = []
	# Assume the file name contains date time.
	times = []
	for file in sorted(glob(files)):
		tmp = os.path.basename(file)
		time = pendulum.datetime(year=int(tmp[14:18]), month=int(tmp[42:44]), day=int(tmp[44:46]), hour=int(tmp[46:48]))
		# f = xr.open_dataset(file, engine='cfgrib', backend_kwargs={
		# 	'filter_by_keys': {
		# 		'typeOfLevel': 'surface',
		# 		'edition': 1
		# 	},
		# 	'indexpath': ''
		# })
		# time = pendulum.parse(str(f.valid_time.values))
		if start_time <= time <= end_time:
			times.append(time)
			used_files.append(file)
	if len(used_files) == 0:
		cli.error(f'Cannot find valid ECWF files!')
	f = xr.open_mfdataset(used_files, engine='cfgrib', concat_dim='valid_time', combine='nested', backend_kwargs={
		'filter_by_keys': {
			'typeOfLevel': 'surface',
			'edition': 1
		},
		'indexpath': ''
	})

	u10  = f.u10.interp(longitude=lon, latitude=lat, method='linear') * units.mps
	v10  = f.v10.interp(longitude=lon, latitude=lat, method='linear') * units.mps
	t2   = f.t2m.interp(longitude=lon, latitude=lat, method='linear') * units.K
	d2   = f.d2m.interp(longitude=lon, latitude=lat, method='linear') * units.K
	ps   = f.sp .interp(longitude=lon, latitude=lat, method='linear') * units.Pa
	slp  = f.msl.interp(longitude=lon, latitude=lat, method='linear') * units.Pa
	sh2  = calc.specific_humidity_from_dewpoint(ps, d2)
	mr2  = calc.mixing_ratio_from_specific_humidity(sh2)
	# tv   = float(calc.virtual_temperature_from_dewpoint(ps, t2, d2))
	tv   = calc.virtual_temperature(t2, mr2)
	gam  = 0.0065
	rd   = 287.05
	g    = 9.80665
	zs   = float(tv / gam * ((slp / ps)**(gam * rd / g) - 1))
	ws10 = calc.wind_speed(u10, v10)
	wd10 = calc.wind_direction(u10, v10)
	rh2  = calc.relative_humidity_from_dewpoint(t2, d2) * 100 * units.percent
	data = {
		'sid' : [sid for _ in range(len(times))],
		'lon' : [lon for _ in range(len(times))],
		'lat' : [lat for _ in range(len(times))],
		'zs'  : [zs],
		'time': times,
		'u10' : [float(u10)],
		'v10' : [float(v10)],
		't2'  : [float(t2)],
		'ps'  : [float(ps)],
		'slp' : [float(slp)],
		'ws10': [float(ws10)],
		'wd10': [float(wd10)],
		'sh2' : [float(sh2)],
		'rh2' : [float(rh2)],
		'd2'  : [float(d2)]
	}
	return pd.DataFrame(data)

def sim_synop_wrf(file_root, sid, lon, lat, start_time, end_time):
	cli.notice(f'Simulate SYNOP at {sid}@({lon},{lat}).')
	files = file_root + '/wrfout*'
	used_files = []
	# Assume the file name contains date time.
	times = []
	for file in sorted(glob(files)):
		f = xr.open_dataset(file, engine='netcdf4')
		time = pendulum.from_format(str(f.Times[0].data)[2:-1], 'YYYY-MM-DD_HH:mm:SS')
		if start_time <= time <= end_time:
			times.append(time)
			used_files.append(file)
	if len(used_files) == 0:
		cli.error(f'Cannot find valid wrfout files!')
	f = xr.open_mfdataset(used_files, concat_dim='Time', combine='nested')

	xlon = f.XLONG[0,:,:]
	xlat = f.XLAT [0,:,:]
	kdtree = create_kdtree(xlon, xlat)
	bbox = find_bbox(kdtree, lon, lat)
	wgts = create_bilinear_wgts(xlon, xlat, bbox, lon, lat)

	zs   = apply_bilinear(wgts, f.HGT , dims=('south_north','west_east')) * units.m
	u10  = apply_bilinear(wgts, f.U10 , dims=('south_north','west_east')) * units.mps
	v10  = apply_bilinear(wgts, f.V10 , dims=('south_north','west_east')) * units.mps
	t2   = apply_bilinear(wgts, f.T2  , dims=('south_north','west_east')) * units.K
	sh2  = apply_bilinear(wgts, f.Q2  , dims=('south_north','west_east'))
	ps   = apply_bilinear(wgts, f.PSFC, dims=('south_north','west_east')) * units.Pa
	slp  = ps * (1 + 0.0065 * units.K / units.m * zs / t2) ** 5.257
	ws10 = calc.wind_speed(u10, v10)
	wd10 = calc.wind_direction(u10, v10)
	rh2  = calc.relative_humidity_from_mixing_ratio(ps, t2, sh2).to(units.percent)
	d2   = calc.dewpoint_from_relative_humidity(t2, rh2).to(units.K)
	data = {
		'sid' : [sid for _ in range(len(times))],
		'lon' : [lon for _ in range(len(times))],
		'lat' : [lat for _ in range(len(times))],
		'zs'  : list(zs.data),
		'time': times,
		'u10' : list(u10.data),
		'v10' : list(v10.data),
		't2'  : list(t2.data),
		'sh2' : list(sh2.data),
		'ps'  : list(ps.data),
		'slp' : list(slp.data),
		'ws10': list(ws10.data),
		'wd10': list(wd10.data),
		'rh2' : list(rh2.data),
		'd2'  : list(d2.data)
	}
	return pd.DataFrame(data)

def add_obs_err(df, obs_err_u, obs_err_v, obs_err_t, obs_err_q, obs_err_p):
	for i in range(len(df)):
		df.loc[i,'u10'] += np.random.normal(loc=0, scale=obs_err_u, size=1)
		df.loc[i,'v10'] += np.random.normal(loc=0, scale=obs_err_v, size=1)
		df.loc[i,'t2' ] += np.random.normal(loc=0, scale=obs_err_t, size=1)
		df.loc[i,'ps' ] += np.random.normal(loc=0, scale=obs_err_p, size=1)
		df.loc[i,'slp'] = df.loc[i,'ps'] * (1 + 0.0065 * df.loc[i,'zs'] / df.loc[i,'t2']) ** 5.257
		sh0 = df.loc[i,'sh2']
		while True:
			sh = sh0 + np.random.normal(loc=0, scale=obs_err_q, size=1)[0]
			if sh < 0: continue
			rh = calc.relative_humidity_from_specific_humidity(df.loc[i,'ps'] * units.Pa, df.loc[i,'t2'] * units.K, sh)
			if rh < 0 or rh > 1: continue
			break
	df['d2'  ] = calc.dewpoint_from_relative_humidity(list(df['t2']) * units.K, list(df['rh2']) * units.percent).to(units.K).magnitude
	df['ws10'] = calc.wind_speed(list(df['u10']) * units.mps, list(df['v10']) * units.mps)
	df['wd10'] = calc.wind_direction(list(df['u10']) * units.mps, list(df['v10']) * units.mps)
	return df

def write_synop_littler(file_path, df):
	f = open(file_path, 'w')
	for idx, row in df.iterrows():
		# Header
		f.write('{:20.5f}' .format(row['lat']))                           # latitude
		f.write('{:20.5f}' .format(row['lon']))                           # longitude
		f.write('{:40}'    .format(row['sid']))                           # id
		f.write('{:40}'    .format(row['sid']))                           # name
		f.write('{:40}'    .format('FM-12 SYNOP'))                        # platform
		f.write('{:40}'    .format('OSSE'))                               # source
		f.write('{:>20.5f}'.format(row['zs']))                            # elevation
		f.write('{:>10d}'  .format(1))                                    # num_vld_fld
		f.write('{:>10d}'  .format(0))                                    # num_error
		f.write('{:>10d}'  .format(0))                                    # num_warning
		f.write('{:>10d}'  .format(idx))                                  # seq_num
		f.write('{:>10d}'  .format(0))                                    # num_dups
		f.write('{:>10}'   .format('F'))                                  # is_sound
		f.write('{:>10}'   .format('F'))                                  # bogus
		f.write('{:>10}'   .format('F'))                                  # discard
		f.write('{:>10d}'  .format(-888888))                              # obs_time
		f.write('{:>10d}'  .format(-888888))                              # julian_day
		f.write('{:>20}'   .format(row['time'].strftime('%Y%m%d%H%M%S'))) # date time
		f.write('{:>13.5f}'.format(row['slp']))                           # slp
		f.write('{:>7d}'   .format(0))                                    # slp qc
		f.write('{:>13.5f}'.format(-888888.0))                            # ref_pres
		f.write('{:>7d}'   .format(0))                                    # ref_pres qc
		f.write('{:>13.5f}'.format(-888888.0))                            # ground_t
		f.write('{:>7d}'   .format(0))                                    # ground_t qc
		f.write('{:>13.5f}'.format(-888888.0))                            # sst
		f.write('{:>7d}'   .format(0))                                    # sst qc
		f.write('{:>13.5f}'.format(-888888.0))                            # psfc
		f.write('{:>7d}'   .format(0))                                    # psfc qc
		f.write('{:>13.5f}'.format(-888888.0))                            # precip
		f.write('{:>7d}'   .format(0))                                    # precip qc
		f.write('{:>13.5f}'.format(-888888.0))                            # t_max
		f.write('{:>7d}'   .format(0))                                    # t_max qc
		f.write('{:>13.5f}'.format(-888888.0))                            # t_min
		f.write('{:>7d}'   .format(0))                                    # t_min qc
		f.write('{:>13.5f}'.format(-888888.0))                            # t_max night
		f.write('{:>7d}'   .format(0))                                    # t_max night qc
		f.write('{:>13.5f}'.format(-888888.0))                            # p_tend03
		f.write('{:>7d}'   .format(0))                                    # p_tend03 qc
		f.write('{:>13.5f}'.format(-888888.0))                            # p_tend24
		f.write('{:>7d}'   .format(0))                                    # p_tend24 qc
		f.write('{:>13.5f}'.format(-888888.0))                            # cloud_cover
		f.write('{:>7d}'   .format(0))                                    # cloud_cover qc
		f.write('{:>13.5f}'.format(-888888.0))                            # ceiling
		f.write('{:>7d}'   .format(0))                                    # ceiling qc
		f.write('\n')
		# Record
		f.write('{:13.5f}'.format(row['ps']))                            # pressure (Pa)
		f.write('{:7d}'   .format(0))                                    # pressure qc
		f.write('{:13.5f}'.format(row['zs']))                            # height
		f.write('{:7d}'   .format(0))                                    # height qc
		f.write('{:13.5f}'.format(row['t2']))                            # temperature (K)
		f.write('{:7d}'   .format(0))                                    # temperature qc
		f.write('{:13.5f}'.format(row['d2']))                            # dewpoint (K)
		f.write('{:7d}'   .format(0))                                    # dewpoint qc
		f.write('{:13.5f}'.format(row['ws10']))                          # wind speed (m s-1)
		f.write('{:7d}'   .format(0))                                    # wind speed qc
		f.write('{:13.5f}'.format(row['wd10']))                          # wind direction (deg)
		f.write('{:7d}'   .format(0))                                    # wind direction qc
		f.write('{:13.5f}'.format(row['u10']))                           # wind u component (m s-1)
		f.write('{:7d}'   .format(0))                                    # wind u component qc
		f.write('{:13.5f}'.format(row['v10']))                           # wind v component (m s-1)
		f.write('{:7d}'   .format(0))                                    # wind v component qc
		f.write('{:13.5f}'.format(row['rh2']))                           # relative humidity (%)
		f.write('{:7d}'   .format(0))                                    # relative humidity qc
		f.write('{:13.5f}'.format(-888888.0))                            # thickness (m)
		f.write('{:7d}'   .format(0))                                    # thickness qc
		f.write('\n')
		# End
		f.write('{:13.5f}'.format(-777777.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-777777.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('\n')
		f.write('{:7d}{:7d}{:7d}\n'.format(1, 0, 0))
	f.close()

def sim_synop(nature_run_root, nature_run_type,
	sid, lon, lat, start_time, end_time,
	obs_err_u, obs_err_v, obs_err_t, obs_err_q, obs_err_p, ntasks,
	output_littler):
	# Check directory existence.
	if not os.path.isdir(os.path.dirname(os.path.abspath(output_littler))):
		cli.notice(f'Create directory {os.path.dirname(output_littler)}.')
		os.makedirs(os.path.dirname(os.path.abspath(output_littler)))

	# Do jobs in parallel.
	n = len(sid)
	pool = Pool(processes=ntasks)
	if nature_run_type == 'wrf':
		sim_synop_call = sim_synop_wrf
	elif nature_run_type == 'era5':
		sim_synop_call = sim_synop_era5
	elif nature_run_type == 'ecwf':
		sim_synop_call = sim_synop_ecwf
	res = pool.starmap(sim_synop_call,
		zip_longest(
			repeat(nature_run_root, n),
			sid, lon, lat,
			repeat(start_time, n),
			repeat(end_time, n)
		)
	)

	# Concatenate results.
	df = []
	for i in range(len(res)):
		df.append(res[i])
	df = pd.concat(df, ignore_index=True)

	# Add observation errors.
	df = add_obs_err(df, obs_err_u, obs_err_v, obs_err_t, obs_err_q, obs_err_p)

	# Write LITTLE_R file.
	write_synop_littler(output_littler, df)

if __name__ == '__main__':
	parser = argparse.ArgumentParser('Simulate SYNOP observation.')
	parser.add_argument('-i', '--nature-run-root')
	parser.add_argument('-t', '--nature-run-type', choices=('ecwf', 'wrf'), required=True)
	parser.add_argument('-o', '--output-littler', required=True)
	parser.add_argument('--sid', nargs='+')
	parser.add_argument('--lon', nargs='+', type=float)
	parser.add_argument('--lat', nargs='+', type=float)
	parser.add_argument('--start-time', required=True, type=parse_time)
	parser.add_argument('--end-time', required=True, type=parse_time)
	parser.add_argument('--obs-err-u', type=float)
	parser.add_argument('--obs-err-v', type=float)
	parser.add_argument('--obs-err-t', type=float)
	parser.add_argument('--obs-err-q', type=float)
	parser.add_argument('--obs-err-p', type=float)
	parser.add_argument('-j', '--json', help='JSON configuration file')
	parser.add_argument('-n', '--ntasks', help='Process number', type=int, default=4)
	args = parser.parse_args()

	if args.json:
		config = parse_config(args.json)
		if not args.nature_run_root:
			args.nature_run_root = config['custom']['nature_run']['dir_pattern']
		if not args.nature_run_type:
			args.nature_run_type = config['custom']['nature_run']['type']
		loc = np.array(config['custom']['sim_obs']['synop']['location'])
		args.sid = [str(int(x)).zfill(5) for x in np.linspace(1, len(loc), len(loc))]
		args.lon = loc[:,0]
		args.lat = loc[:,1]
		args.obs_err_u = config['custom']['sim_obs']['synop']['obs_err']['u']
		args.obs_err_v = config['custom']['sim_obs']['synop']['obs_err']['v']
		args.obs_err_t = config['custom']['sim_obs']['synop']['obs_err']['t']
		args.obs_err_q = config['custom']['sim_obs']['synop']['obs_err']['q']
		args.obs_err_p = config['custom']['sim_obs']['synop']['obs_err']['p']

	sim_synop(args.nature_run_root, args.nature_run_type,
		args.sid, args.lon, args.lat, args.start_time, args.end_time,
		args.obs_err_u, args.obs_err_v, args.obs_err_t, args.obs_err_q, args.obs_err_p, args.ntasks,
		args.output_littler)
